﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using zlib;

namespace dnAnalytics.LinearAlgebra.IO.Matlab
{
    internal class MatlabParser
    {
        private const int _largeBlockSize = 8;
        private const byte _littleEndianIndicator = 0x49;
        private const int _smallBlockSize = 4;
        private readonly IList<string> _names = new List<string>();
        private readonly Stream _stream;


        public MatlabParser(string fileName)
            : this(fileName, new string[0])
        {
        }

        public MatlabParser(Stream stream)
            : this(stream, new string[0])
        {
        }

        public MatlabParser(Stream stream, IEnumerable<string> objectNames)
        {
            if (stream == null)
            {
                throw new ArgumentNullException("stream");
            }

            _stream = stream;
            SetNames(objectNames);
        }

        public MatlabParser(string fileName, IEnumerable<string> objectNames)
        {
            if (string.IsNullOrEmpty(fileName))
            {
                throw new ArgumentException("File name cannot be null or empty.", "fileName");
            }
            _stream = File.OpenRead(fileName);
            SetNames(objectNames);
        }

        private void SetNames(IEnumerable<string> objectNames)
        {
            foreach (string name in objectNames)
            {
                _names.Add(name);
            }
        }

        public MatlabFile Parse()
        {
            MatlabFile file = new MatlabFile();

            using (BinaryReader reader = new BinaryReader(_stream))
            {
                file.HeaderText = Encoding.ASCII.GetString(reader.ReadBytes(116));

                //skipping subsystem offsets
                reader.BaseStream.Position = 126;

                if (reader.ReadByte() != _littleEndianIndicator)
                {
                    throw new NotSupportedException("Big endian files are not supported.");
                }

                //skip version since it is always 0x0100.
                reader.BaseStream.Position = 128;
                long length = _stream.Length;

                // for each data block add a matlab object to the file.
                while (reader.BaseStream.Position < length)
                {
                    DataType type = (DataType) reader.ReadInt16();
                    int size = reader.ReadInt16();
                    bool smallBlock = true;
                    if (size == 0)
                    {
                        size = reader.ReadInt32();
                        smallBlock = false;
                    }

                    byte[] data;
                    if (type == DataType.Compressed)
                    {
                        data = DecompressBlock(reader.ReadBytes(size), ref type);
                    }
                    else
                    {
                        data = new byte[size];
                        reader.Read(data, 0, size);
                        AlignData(reader.BaseStream, size, smallBlock);
                    }

                    if (type == DataType.Matrix)
                    {
                        AddMatrix(data, file);
                    }
                    else
                    {
                        throw new NotSupportedException(string.Format("{0} is not a supported type.", type));
                    }
                }
            }
            return file;
        }


        private static void AlignData(Stream stream, int size, bool smallBlock)
        {
            int blockSize = smallBlock ? _smallBlockSize : _largeBlockSize;
            int offset = 0;
            int mod = size%blockSize;
            if (mod != 0)
            {
                offset = blockSize - mod;
            }

            stream.Seek(offset, SeekOrigin.Current);
        }

        private static byte[] DecompressBlock(byte[] compressed, ref DataType type)
        {
            byte[] data;
            using(MemoryStream decompressed = new MemoryStream())
            using (ZOutputStream decompressor = new ZOutputStream(decompressed))
            {
                decompressor.Write(compressed, 0, compressed.Length);
                decompressed.Position = 0;
                byte[] buf = new byte[4];
                decompressed.Read(buf, 0, 4);
                type = (DataType) BitConverter.ToInt32(buf, 0);
                decompressed.Read(buf, 0, 4);
                int size = BitConverter.ToInt32(buf, 0);
                data = new byte[size];
                decompressed.Read(data, 0, size);
            }
            return data;
        }

        private void AddMatrix(byte[] data, MatlabFile file)
        {
            using (MemoryStream ms = new MemoryStream(data))
            using (BinaryReader reader = new BinaryReader(ms))
            {
                //skip tag - doesn't tell us anything we don't already know
                reader.BaseStream.Seek(8, SeekOrigin.Current);

                ArrayClass arrayClass = (ArrayClass) reader.ReadByte();
                byte flags = reader.ReadByte();
                bool isComplex = (flags & (byte) ArrayFlags.Complex) == (byte) ArrayFlags.Complex;

                if( isComplex)
                {
                    throw new NotSupportedException(Properties.Resources.ComplexNotSupported);
                }
                //skip unneeded bytes
                reader.BaseStream.Seek(10, SeekOrigin.Current);

                int numDimensions = reader.ReadInt32()/8;
                if (numDimensions > 2)
                {
                    throw new NotSupportedException(Properties.Resources.MoreThan2D);
                }
                int rows = reader.ReadInt32();
                int columns = reader.ReadInt32();

                //skip unneeded bytes
                reader.BaseStream.Seek(2, SeekOrigin.Current);
                int size = reader.ReadInt16();
                bool smallBlock = true;
                if (size == 0)
                {
                    size = reader.ReadInt32();
                    smallBlock = false;
                }

                string name = Encoding.ASCII.GetString(reader.ReadBytes(size));
                AlignData(reader.BaseStream, size, smallBlock);

                //only grab wanted objects
                if (_names.Count != 0 && !_names.Contains(name))
                {
                    return;
                }

                DataType type = (DataType) reader.ReadInt16();
                size = reader.ReadInt16();
                if (size == 0)
                {
                    size = reader.ReadInt32();
                }

                Matrix matrix;
                if (arrayClass == ArrayClass.Sparse)
                {
                    matrix = PopulateSparseMatrix(reader, rows, columns, size);
                }
                else if (arrayClass == ArrayClass.Unknown || arrayClass == ArrayClass.Cell || arrayClass == ArrayClass.Structure
                         || arrayClass == ArrayClass.Object || arrayClass == ArrayClass.Character || arrayClass == ArrayClass.Function)
                {
                    throw new NotImplementedException();
                }
                else
                {
                    matrix = PopulateDenseMatrix(type, reader, rows, columns);
                }
                file.Matrices.Add(name, matrix);
                if (file.FirstMatrixName == null)
                {
                    file.FirstMatrixName = name;
                }
            }
        }

        private static Matrix PopulateSparseMatrix(BinaryReader reader, int rows, int columns, int size)
        {
            //populate the row data array
            int[] ir = new int[size/4];
            for (int i = 0; i < ir.Length; i++)
            {
                ir[i] = reader.ReadInt32();
            }
            AlignData(reader.BaseStream, size, false);
            
            //skip data type since it will always be int32
            reader.BaseStream.Seek(4, SeekOrigin.Current);

            //populate the column data array
            int jcSize = reader.ReadInt32();
            int[] jc = new int[jcSize/4];
            for (int j = 0; j < jc.Length; j++)
            {
                jc[j] = reader.ReadInt32();
            }
            AlignData(reader.BaseStream, jcSize, false);

            DataType type = (DataType)reader.ReadInt32();

            //skip length since we already no it for the number of rows
            reader.BaseStream.Seek(4, SeekOrigin.Current);

            Matrix matrix = new SparseMatrix(rows, columns);
            int col = 0;
            for (int i = 0; i < ir.Length; i++)
            {
                int row = ir[i];
                if (jc[col+1] == i)
                {
                    col++;
                }
                switch (type)
                {
                    case DataType.Int8:
                        matrix[row, col] = reader.ReadSByte();
                        break;
                    case DataType.UInt8:
                        matrix[row, col] = reader.ReadByte();
                        break;
                    case DataType.Int16:
                        matrix[row, col] = reader.ReadInt16();
                        break;

                    case DataType.UInt16:
                        matrix[row, col] = reader.ReadUInt16();
                        break;
                    case DataType.Int32:
                        matrix[row, col] = reader.ReadInt32();
                        break;

                    case DataType.UInt32:
                        matrix[row, col] = reader.ReadUInt32();
                        break;
                    case DataType.Single:
                        matrix[row, col] = reader.ReadSingle();
                        break;
                    case DataType.Int64:
                        matrix[row, col] = reader.ReadInt64();
                        break;
                    case DataType.UInt64:
                        matrix[row, col] = reader.ReadUInt64();
                        break;
                    case DataType.Double:
                        matrix[row, col] = reader.ReadDouble();
                        break;
                    default:
                        throw new NotSupportedException();
                }
            }
            return matrix;
        }

        private static Matrix PopulateDenseMatrix(DataType type, BinaryReader reader, int rows, int columns)
        {
            Matrix matrix = new DenseMatrix(rows, columns);
            switch (type)
            {
                case DataType.Int8:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadSByte();
                        }
                    }
                    break;
                case DataType.UInt8:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadByte();
                        }
                    }
                    break;
                case DataType.Int16:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadInt16();
                        }
                    }
                    break;
                case DataType.UInt16:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadUInt16();
                        }
                    }
                    break;
                case DataType.Int32:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadInt32();
                        }
                    }
                    break;
                case DataType.UInt32:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadUInt32();
                        }
                    }
                    break;
                case DataType.Single:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadSingle();
                        }
                    }
                    break;
                case DataType.Int64:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadInt64();
                        }
                    }
                    break;
                case DataType.UInt64:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadUInt64();
                        }
                    }
                    break;
                case DataType.Double:
                    for (int j = 0; j < columns; j++)
                    {
                        for (int i = 0; i < rows; i++)
                        {
                            matrix[i, j] = reader.ReadDouble();
                        }
                    }
                    break;
                default:
                    throw new NotSupportedException();
            }

            return matrix;
        }
    }
}